from time import time
import multiprocessing as mp
import torch
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader, TensorDataset

resize=224
data_transform = [transforms.ToTensor()]
if resize:
    data_transform.insert(0, transforms.Resize(resize))
data_compose = transforms.Compose(data_transform)  # 数据预处理为tensor类型
cifar10_train = torchvision.datasets.CIFAR10(
    root="./dataset", train=True, transform=data_compose, download=True
)


print(f"num of CPU: {mp.cpu_count()}")
for num_workers in range(2, mp.cpu_count(), 2):  
    train_loader = DataLoader(
    dataset=cifar10_train, batch_size=64, num_workers=num_workers, shuffle=True, drop_last=False
    )
    start = time()
    for epoch in range(1, 3):
        for i, data in enumerate(train_loader, 0):
            pass
    end = time()
    print("Finish with:{} second, num_workers={}".format(end - start, num_workers))
